{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9e522389-2890-42f9-a1af-9582136d118d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dataclasses\n",
    "import functools\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import jax\n",
    "import haiku as hk\n",
    "\n",
    "from graphcast import graphcast\n",
    "from graphcast import autoregressive\n",
    "from graphcast import casting\n",
    "from graphcast import normalization\n",
    "from graphcast import rollout\n",
    "from graphcast import data_utils\n",
    "from graphcast import checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dedcb2ff-f628-4904-b8b3-e6130df9683f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The weights and norm files can be downloaded from Google's official Cloud Storage.\n",
    "\n",
    "with open(\"params/graphcast_params_GraphCast - ERA5 1979-2017 - resolution 0.25 - pressure levels 37 - mesh 2to6 - precipitation input and output.npz\", \"rb\") as f:\n",
    "    ckpt = checkpoint.load(f, graphcast.CheckPoint)\n",
    "\n",
    "params = ckpt.params\n",
    "model_config = ckpt.model_config\n",
    "task_config = ckpt.task_config\n",
    "\n",
    "state = {}\n",
    "\n",
    "with open(\"stats/graphcast_stats_diffs_stddev_by_level.nc\", \"rb\") as f:\n",
    "    diffs_stddev_by_level = xr.load_dataset(f).compute()\n",
    "\n",
    "with open(\"stats/graphcast_stats_mean_by_level.nc\", \"rb\") as f:\n",
    "    mean_by_level = xr.load_dataset(f).compute()\n",
    "\n",
    "with open(\"stats/graphcast_stats_stddev_by_level.nc\", \"rb\") as f:\n",
    "    stddev_by_level = xr.load_dataset(f).compute()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8a5cfa2b-0128-4f4e-a4ed-1260cbbe8461",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ModelConfig(resolution=0.25, mesh_size=6, latent_size=512, gnn_msg_steps=16, hidden_layers=1, radius_query_fraction_edge_length=0.5999912857713345, mesh2grid_edge_normalization_factor=0.6180338738074472)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f5f68b23-a6e8-464e-bfd5-4c299eb8db16",
   "metadata": {},
   "outputs": [],
   "source": [
    "def construct_wrapped_graphcast(\n",
    "        model_config: graphcast.ModelConfig,\n",
    "        task_config: graphcast.TaskConfig\n",
    "        ):\n",
    "    \n",
    "    predictor = graphcast.GraphCast(model_config, task_config)\n",
    "\n",
    "    predictor = casting.Bfloat16Cast(predictor)\n",
    "\n",
    "    predictor = normalization.InputsAndResiduals(\n",
    "        predictor,\n",
    "        diffs_stddev_by_level=diffs_stddev_by_level,\n",
    "        mean_by_level=mean_by_level,\n",
    "        stddev_by_level=stddev_by_level\n",
    "    )\n",
    "\n",
    "    predictor = autoregressive.Predictor(predictor, gradient_checkpointing=True)\n",
    "    \n",
    "    return predictor\n",
    "\n",
    "\n",
    "@hk.transform_with_state\n",
    "def run_forward(model_config, task_config, inputs, targets_template, forcings):\n",
    "    predictor = construct_wrapped_graphcast(model_config, task_config)\n",
    "    return predictor(inputs, targets_template=targets_template, forcings=forcings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3c7786e7-ce9b-408a-9c6f-4efb32e513c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def with_configs(fn):\n",
    "    return functools.partial(\n",
    "        fn, model_config=model_config, task_config=task_config)\n",
    "\n",
    "def with_params(fn):\n",
    "    return functools.partial(fn, params=params, state=state)\n",
    "\n",
    "def drop_state(fn):\n",
    "    return lambda **kw: fn(**kw)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "86d3adaa-4806-4752-b3cd-cf0c97ea884d",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_jitted = jax.jit(with_configs(run_forward.init))\n",
    "\n",
    "run_forward_jitted = drop_state(with_params(jax.jit(with_configs(run_forward.apply))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cb0acfb0-7126-48f1-8e11-1d5aee3b3361",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_data = xr.open_dataset(\"data/201905.nc\", engine='netcdf4')\n",
    "example_data = example_data.drop_vars(['batch', 'pressure_level', 'number', 'expver']) # Optional, there is also a demo about constructing example in the current directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "131c53a8-aa81-4549-b4ce-9a5804a8a9c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_inputs, eval_targets, eval_forcings = data_utils.extract_inputs_targets_forcings(\n",
    "    example_data,\n",
    "    target_lead_times=slice(\"6h\", \"6h\"),\n",
    "    **dataclasses.asdict(task_config)\n",
    ")\n",
    "rng = jax.random.PRNGKey(42)\n",
    "rollout_targets_template = eval_targets * np.nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "86b9b8e5-6b7a-4c12-a42b-72104950b5b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "tot_steps = 20\n",
    "multi_step_rollout_targets_template = rollout.extend_targets_template(rollout_targets_template, tot_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "298a8b50-2313-410e-9d05-7354b7ed180c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def construct_full_forcings(\n",
    "    example_data: xr.Dataset,\n",
    "    start_lead_time: str = \"6h\",\n",
    "    num_steps: int = 20,\n",
    "    step_hours: int = 6,\n",
    "    forcing_variables: tuple = (\n",
    "        \"day_progress_sin\", \"day_progress_cos\",\n",
    "        \"year_progress_sin\", \"year_progress_cos\",\n",
    "        \"toa_incident_solar_radiation\"\n",
    "    )\n",
    ") -> xr.Dataset:\n",
    "\n",
    "    datetime_coord = example_data.coords[\"datetime\"]\n",
    "    \n",
    "    if \"batch\" in datetime_coord.dims:\n",
    "        base_datetime = datetime_coord.isel(batch=0).values[-1]\n",
    "    else:\n",
    "        base_datetime = datetime_coord.values[-1]\n",
    "\n",
    "    start_hours = int(start_lead_time.rstrip(\"h\"))\n",
    "\n",
    "    datetimes = [\n",
    "        base_datetime + np.timedelta64(start_hours + i * step_hours, 'h')\n",
    "        for i in range(num_steps)\n",
    "    ]\n",
    "    lead_times = [f\"{start_hours + i * step_hours}h\" for i in range(num_steps)]\n",
    "    datetimes = np.array(datetimes, dtype=\"datetime64[ns]\")  # shape: (num_steps,)\n",
    "\n",
    "    coords = {\n",
    "        \"time\": lead_times,\n",
    "        \"datetime\": (\"time\", datetimes),\n",
    "        \"lat\": example_data.coords[\"lat\"],\n",
    "        \"lon\": example_data.coords[\"lon\"]\n",
    "    }\n",
    "    forcing_data = xr.Dataset(coords=coords)\n",
    "\n",
    "    if set(forcing_variables) & data_utils._DERIVED_VARS:\n",
    "        data_utils.add_derived_vars(forcing_data)\n",
    "\n",
    "    if set(forcing_variables) & {data_utils.TISR}:\n",
    "        data_utils.add_tisr_var(forcing_data)\n",
    "\n",
    "    forcing_data = forcing_data[list(forcing_variables)]\n",
    "\n",
    "    if \"batch\" in example_data.dims:\n",
    "        if \"batch\" not in forcing_data.dims:\n",
    "            forcing_data = forcing_data.expand_dims(\"batch\", axis=0)\n",
    "        else:\n",
    "            print(\"Warning: 'batch' dim already exists in forcings. Skipping expand_dims.\")\n",
    "\n",
    "        if \"batch\" not in forcing_data.coords[\"datetime\"].dims:\n",
    "            old_dt = forcing_data.coords[\"datetime\"].values  # shape: (time,)\n",
    "            forcing_data = forcing_data.assign_coords({\n",
    "                \"datetime\": ((\"batch\", \"time\"), old_dt[np.newaxis, :])\n",
    "            })\n",
    "\n",
    "    return forcing_data\n",
    "\n",
    "full_forcings = construct_full_forcings(\n",
    "    example_data=example_data,\n",
    "    start_lead_time=\"6h\",\n",
    "    num_steps=tot_steps,\n",
    "    forcing_variables=task_config.forcing_variables,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8e351e0d-0e58-4afe-b2a9-eb6c621843cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Skipping gradient checkpointing for sequence length of 1\n"
     ]
    }
   ],
   "source": [
    "# Multistep inference on A800 80GB, or manually rollout single-step on e.g. >= 4090 24GB\n",
    "predictions = rollout.chunked_prediction(\n",
    "    run_forward_jitted,\n",
    "    rng=rng,\n",
    "    inputs=eval_inputs,\n",
    "    targets_template=multi_step_rollout_targets_template,\n",
    "    forcings=full_forcings,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de7dedf9-a709-46b3-9150-8f710bcd99fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Monitor GPU usage using 'nvitop' or 'nvidia-smi' in the terminal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "07080fa4-4f17-4b4f-84e4-c3ef0bc04133",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FrozenMappingWarningOnValuesAccess({'time': 20, 'batch': 1, 'lat': 721, 'lon': 1440, 'level': 37})"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predictions.dims"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
